from torchair._ge_concrete_graph.ge_converter.converter_utils import *


@declare_supported([
    Support(F32(2, 2), F32(2, 2), 0.),
    Support(F32(2, 2), F32(2, 2), 0.1),
    Support(F16(2, 2), F16(2, 2), 0.),
    Support(F16(2, 2), F16(2, 2), 0.1),
])
@register_fx_node_ge_converter(torch.ops.aten.threshold_backward.default)
def conveter_aten_threshold_backward_default(
    grad_output: Tensor,
    self: Tensor,
    threshold: Union[Number, Tensor],
    meta_outputs: TensorSpec = None,
):
    """ NB: aten::threshold_backward(Tensor grad_output, Tensor self, Scalar threshold) -> Tensor """
    self, grad_output = dtype_promote(self, grad_output, target_dtype=meta_outputs.dtype)

    # This scenario will be supported in the future.
    if isinstance(threshold, Tensor):
        raise RuntimeError("torch.ops.aten.threshold_backward is not implemented while threshold is Tensor!")

    # Suggest a high-performance implementation when threshold is euqal to zero.
    if threshold == 0.:
        return ge.ReluGrad(grad_output, self)
    
    return ge.ThresholdGradV2D(grad_output, self, threshold=threshold)


@register_fx_node_ge_converter(torch.ops.aten.threshold_backward.grad_input)
def conveter_aten_threshold_backward_grad_input(
    grad_output: Tensor,
    self: Tensor,
    threshold: Union[Number, Tensor],
    *,
    grad_input: Tensor = None,
    meta_outputs: TensorSpec = None
):
    """NB: aten::threshold_backward.grad_input(Tensor grad_output, Tensor self, Scalar threshold, *, Tensor(a!) grad_input) -> Tensor(a!)"""
    raise RuntimeError("torch.ops.aten.threshold_backward.grad_input ge_converter is "
                       "redundant before pytorch 2.1.0, might be supported in future version.")
